package com.fliad.ai.service;

import com.fliad.ai.entity.AiDocument;
import com.fliad.ai.entity.AiDocumentChunk;
import com.fliad.ai.entity.AiKnowledge;
import com.fliad.ai.entity.AiLlm;
import org.noear.snack.ONode;
import org.noear.solon.ai.embedding.EmbeddingModel;
import org.noear.solon.ai.rag.Document;
import org.noear.solon.ai.rag.RepositoryStorable;
import org.noear.solon.ai.rag.repository.InMemoryRepository;
import org.noear.solon.ai.rag.splitter.RegexTextSplitter;
import org.noear.solon.ai.rag.splitter.SplitterPipeline;
import org.noear.solon.ai.rag.util.QueryCondition;
import org.noear.solon.annotation.Component;
import org.noear.solon.annotation.Init;
import org.noear.solon.annotation.Inject;
import org.noear.solon.cloud.CloudClient;
import org.noear.solon.cloud.model.Media;
import org.noear.wood.DbContext;
import org.noear.wood.annotation.Db;

import java.io.IOException;
import java.sql.SQLException;
import java.util.List;

import static com.fliad.ai.service.AiDocumentService.KNOWLEDGE_REPO;

@Component
public class AiKnowledgeService {

    @Db
    DbContext db;
    @Inject
    AiDocumentService documentService;

    public boolean create(AiKnowledge knowledge) throws SQLException {
        return db.table("ai_knowledge")
                .setEntity(knowledge)
                .insert() > 0;
    }

    public AiKnowledge get(Long id) throws SQLException {
        return db.table("ai_knowledge")
                .whereEq("id", id)
                .selectItem("*", AiKnowledge.class);
    }

    public List<AiKnowledge> list(String title, String vectorStoreType) throws SQLException {
        return db.table("ai_knowledge")
                .where(buildQueryCondition(title, vectorStoreType))
                .orderBy("id DESC")
                .selectList("*", AiKnowledge.class);
    }

    // 新增方法：构建查询条件
    private String buildQueryCondition(String title, String vectorStoreType) {
        StringBuilder condition = new StringBuilder("1=1");
        if (title != null && !title.isEmpty()) {
            condition.append(" AND title LIKE '%").append(title).append("%'");
        }
        if (vectorStoreType != null && !vectorStoreType.isEmpty()) {
            condition.append(" AND vector_store_type = '").append(vectorStoreType).append("'");
        }
        return condition.toString();
    }

    public boolean update(Long id, AiKnowledge knowledge) throws SQLException {
        return db.table("ai_knowledge")
                .whereEq("id", id)
                .setEntity(knowledge)
                .update() > 0;
    }

    public boolean delete(Long id) throws SQLException, IOException {
        documentService.deleteByKnowledgeId(id);

        return db.table("ai_knowledge")
                .whereEq("id", id)
                .delete() > 0;
    }

    // 新增方法：初始化知识库
    public void initializeKnowledge() throws SQLException, IOException {
        List<AiDocument> aiDocumentList = db.table("ai_document").selectList("*", AiDocument.class);
        for (AiDocument aiDocument : aiDocumentList) {
            ONode options = ONode.loadStr(aiDocument.getOptions());
            SplitterPipeline splitterPipeline = new SplitterPipeline();
            if ("org.noear.solon.ai.rag.splitter.RegexTextSplitter".equals(options.get("splitter").getString())) {
                splitterPipeline.next(new RegexTextSplitter(options.get("regex").getString()));
            }
            Media media = CloudClient.file().get("ai", aiDocument.getFileKey());
            if (media == null || media.contentSize() == 0) {
                continue;
            }
            List<Document> docList = splitterPipeline.split(media.bodyAsString());
            for (int i = 0; i < docList.size(); i++) {
                docList.get(i).metadata("name", aiDocument.getFileKey());
                docList.get(i).metadata("knowledgeId", aiDocument.getKnowledgeId());
                docList.get(i).metadata("documentId", aiDocument.getId());
            }
            AiKnowledge aiKnowledge = db.table("ai_knowledge")
                    .whereEq("id", aiDocument.getKnowledgeId())
                    .selectItem("*", AiKnowledge.class);
            AiLlm aiLlm = db.table("ai_llm")
                    .whereEq("id", aiKnowledge.getVectorEmbedLlmId())
                    .selectItem("*", AiLlm.class);
            RepositoryStorable repositoryStorable = KNOWLEDGE_REPO.computeIfAbsent(aiDocument.getKnowledgeId(), key -> {
                return new InMemoryRepository(EmbeddingModel.of(aiLlm.getLlmEndpoint())
                        .apiKey(aiLlm.getLlmApiKey()).provider(aiLlm.getDialect()).model(aiLlm.getLlmModel()).build());
            });
            List<AiDocumentChunk> aiDocumentChunkList = db.table("ai_document_chunk")
                    .whereEq("document_id", aiDocument.getId())
                    .selectList("*", AiDocumentChunk.class);
            repositoryStorable.delete(aiDocumentChunkList.stream().map(AiDocumentChunk::getId).toArray(String[]::new));
            repositoryStorable.insert(docList);
        }
    }

    public Object search(Long id, String keyword) throws SQLException, IOException {
        AiKnowledge aiKnowledge = db.table("ai_knowledge").whereEq("id", id)
                .selectItem("*", AiKnowledge.class);
        AiLlm aiLlm = db.table("ai_llm")
                .whereEq("id", aiKnowledge.getVectorEmbedLlmId())
                .selectItem("*", AiLlm.class);
        RepositoryStorable repositoryStorable = KNOWLEDGE_REPO.computeIfAbsent(aiKnowledge.getId(), key -> {
            return new InMemoryRepository(EmbeddingModel.of(aiLlm.getLlmEndpoint())
                    .apiKey(aiLlm.getLlmApiKey()).provider(aiLlm.getDialect()).model(aiLlm.getLlmModel()).build());
        });
        QueryCondition condition = new QueryCondition(keyword);
        condition.filterExpression("knowledgeId == " + id + "L");
        return repositoryStorable.search(condition);
    }

}
